import logging
import time
import numpy as np
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from collections import Counter
from pinecone import Pinecone, ServerlessSpec
from tqdm import tqdm

# 配置 logging，带日期
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s %(levelname)s: %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)


# 初始化Pinecone客户端
pc = Pinecone(api_key="pcsk_dqdo1_EP62eXyBSFyv5dQWZiMJZybi6YAkjrT8cYaobXeTkRDvfnbvxY9y8RCdZkajBb1")
index_name = "mnist-index"

# 检查并清除现有索引
existing_indexes = pc.list_indexes()
if any(index['name'] == index_name for index in existing_indexes):
    logging.info(f"删除现有索引 '{index_name}'...")
    pc.delete_index(index_name)

# 创建新索引
logging.info(f"创建索引 '{index_name}'...")
pc.create_index(
    name=index_name,
    dimension=64,
    metric="euclidean",
    spec=ServerlessSpec(cloud="aws", region="us-east-1")
)
# 等待索引 ready
while True:
    desc = pc.describe_index(index_name)
    if desc['status']['ready']:
        break
    time.sleep(2)
index = pc.Index(index_name)
logging.info(f"已连接到索引 '{index_name}'")

# 加载MNIST数据集并划分训练/测试集
X, y = load_digits(n_class=10, return_X_y=True)
X = X.astype(np.float32)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

# 上传训练集到 Pinecone
vectors = [(str(i), X_train[i].tolist(), {"label": int(y_train[i])}) for i in range(len(X_train))]
logging.info(f"开始上传训练集向量，共 {len(vectors)} 条...")
batch_size = 500
for i in tqdm(range(0, len(vectors), batch_size), desc="上传进度"):
    batch = vectors[i:i + batch_size]
    index.upsert(vectors=batch)
logging.info(f"成功创建索引，并上传了{len(vectors)}条数据")

# 查询并评估准确率
k = 11
correct = 0
logging.info(f"开始测试，k={k}，测试集样本数：{len(X_test)}")
for i in tqdm(range(len(X_test)), desc="测试进度"):
    query_vec = X_test[i].tolist()
    results = index.query(vector=query_vec, top_k=k, include_metadata=True)
    if not results['matches']:
        continue
    labels = [match['metadata']['label'] for match in results['matches']]
    pred = Counter(labels).most_common(1)[0][0]
    if pred == y_test[i]:
        correct += 1
accuracy = correct / len(X_test)
logging.info(f"当k={k}时，使用Pinecone的准确率为：{accuracy:.4f}")
